from joblib import load
import numpy as np

class Classification(object):
    """股票分类预测工具类（使用预训练模型）"""
    
    def __init__(self, model_path='knn_classifier.joblib', 
                 scaler_path='feature_scaler.joblib', 
                 encoder_path='label_encoder.joblib'):
        """
        初始化加载预训练模型和预处理工具
        参数:
            model_path: KNN模型路径
            scaler_path: 特征标准化器路径
            encoder_path: 标签编码器路径
        """
        try:
            self.knn = load(model_path)         # 加载KNN分类模型
            self.scaler = load(scaler_path)      # 加载特征标准化器
            self.le = load(encoder_path)         # 加载标签编码器
            print("模型和预处理工具加载成功！")
        except Exception as e:
            raise ValueError(f"模型加载失败: {str(e)}")

    def get_predicted_category(self, new_data):
        """
        预测新数据的股票分类
        参数:
            new_data: 字典格式的输入数据，包含所有必要特征
        返回:
            预测的类别标签（字符串）
        """
        # 定义特征顺序（必须与训练时一致）
        feature_order = [
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ]
        
        # 检查输入数据完整性
        missing_features = set(feature_order) - set(new_data.keys())
        if missing_features:
            raise ValueError(f"缺少必要特征字段: {missing_features}")
        
        try:
            # 转换为二维数组并保持特征顺序
            new_value = np.array([[float(new_data[col]) for col in feature_order]])
            
            # 标准化新数据（使用训练时的scaler）
            new_scaled = self.scaler.transform(new_value)
            
            # 预测分类
            predicted_label = self.knn.predict(new_scaled)
            predicted_category = self.le.inverse_transform(predicted_label)
            
            return predicted_category[0]
        except Exception as e:
            raise ValueError(f"预测过程中出错: {str(e)}")

if __name__ == '__main__':
    # 示例用法
    try:
        # 初始化分类器
        clf = Classification()
        
        # 新数据输入（注意：值需要转换为字符串或数字）
        test_data = {
            'eps': '-0.8309',
            'total_revenue_ps': '11.9673',
            'undist_profit_ps': '7.0209',
            'gross_margin': '11586700000',  # 建议统一使用科学计数法或相同数量级
            'fcff': '-12879100000',
            'fcfe': '-52634600000',
            'tangible_asset': '179873000000',
            'bps': '20.2568',
            'grossprofit_margin': '8.1152',
            'npta': '-0.5821'
        }
        
        # 获取预测结果
        result = clf.get_predicted_category(test_data)
        print(f"\n预测分类结果: {result}")
        
    except Exception as e:
        print(f"程序运行出错: {str(e)}")